package com.fwmagic.spark.ml.naivebayes

import com.fwmagic.spark.util.SparkUtils
import org.apache.spark.ml.classification.{NaiveBayes, NaiveBayesModel}
import org.apache.spark.sql.{DataFrame, SparkSession}

/**
 * 朴素贝叶斯算法
 * 调用Sparkmllib中现成的朴素贝叶斯算法，进行出轨预测（分类）
 *
 * 1.加载样本数据
 * 2.将特征向量化
 * 3.构造朴素贝叶斯算法工具
 * 4.用算法工具对样本训练集训练模型
 * 5.保存模型
 *
 * 样本数据：
 * name,job,income,age,sex,label
 * 张飞,老师,中,青年,男,出轨
 * 赵云,老师,中,中年,女,出轨
 * 陆小凤,老师,低,青年,男,没出
 */
object FitAndSaveModel {

  def main(args: Array[String]): Unit = {
    val spark: SparkSession = SparkUtils.getSparkSession(this.getClass.getSimpleName)

    //1.加载数据
    val sample: DataFrame = spark.read.option("header", true).csv("data/naivebayes/sample.csv")

    import com.fwmagic.spark.ml.utils.VectorUtils._
    spark.udf.register("arr2vec2", arr2vec2)

    //2.将特征向量化（2.1:特征数字化,2.2:数字向量化）
    //2.1:特征数字化
    val sampleVecs: DataFrame = sample.selectExpr(
      "name",
      "case label when '出轨' then 0.0 else 1.0 end as label",
      "arr2vec2(array(job,income,age,sex)) as vec"
    )


    //2.2:数字向量化
    //val sampleVecs: DataFrame = df2.selectExpr("name", "label", "arr2vec(array(cast(job as double), cast(income as double), cast(age as double), cast(sex as double))) as vec")
    //    val sampleVecs: DataFrame = df2.select('name, 'label.cast(DataTypes.DoubleType), arr2vec2(array('job.cast(DataTypes.DoubleType), 'income.cast(DataTypes.DoubleType), 'age.cast(DataTypes.DoubleType), 'sex.cast(DataTypes.DoubleType))).as("vec"))

    //3.构造朴素贝叶斯算法工具
    val bayes = new NaiveBayes()
      .setFeaturesCol("vec")
      .setLabelCol("label")
      .setSmoothing(1.0) //拉普拉斯平滑系数

    //4.用算法工具对样本训练集训练模型
    val model: NaiveBayesModel = bayes.fit(sampleVecs)

    //5.保存模型
    model.write.overwrite.save("data/naivebayes/model")

    println("model is Saved !")

    spark.close()
  }

}
